import sys
import pexpect
import re

DEFAULT_LOGIN_TIMEOUT = 10

class TransportError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class AuthenticationError(TransportError): pass
class NetworkError(TransportError): pass
class LoginError(TransportError): pass

class BaseTransport(object):
    def __init__(self, prompt='', pager_prompt=''):
        self.prompt = prompt
        self.pager_prompt = pager_prompt

        self.last_prompt = ''
 
    def open(self):
        pass

    def close(self):
        self.child.close()

    MATCH_PAGER_PROMPT = 1

    def cmd(self, line):
        output = ''
        self.child.sendline(line)
        while 1:
            result = self.child.expect([self.prompt, self.pager_prompt])
            output +=  self.child.before
            if (result == MATCH_PAGER_PROMPT):
                break
            self.child.send(" ")

        #Update last prompta
        self.last_prompt = self.child.after

        #strip the command from the output
        output = re.sub(line + "\r?\n", '', output, len(line)+2)
        return output


class SSHTransport(BaseTransport):
    def __init__(self, ssh_options='',
                       login_timeout=DEFAULT_LOGIN_TIMEOUT,
                       *args, **kwargs):
        super(SSHTransport, self).__init__(*args, **kwargs) 
        self.ssh_options = ""
        self.login_timeout=login_timeout

    LOGIN_EXPECTS = ["(?i)(?:password)|(?:passphrase for key)",		        # Password:   -		     0
		             "(?i)are you sure you want to continue connecting",	# Accept License 	     1
                     "(?i)permission denied",                               # Bad login              2
                     "(?i)no route to host",                                # No route to host       3
                     "(?i)connection refused",                              # Connection refused     4
                     pexpect.EOF,			                    			# Returned without error 5
                     pexpect.TIMEOUT,				                        # We timed out           6
                    ]
    MATCH_PASSWORD       = 0
    MATCH_ACCEPT_LICENSE = 1
    MATCH_PERM_DENIED    = 2
    MATCH_NO_ROUTE       = 3
    MATCH_CONN_REFUSED   = 4
    MATCH_EOF            = 5
    MATCH_TIMEOUT        = 6
    MATCH_PROMPT         = len(LOGIN_EXPECTS)

    def open(self, host, username="", password="", port=22):
        self.child = pexpect.spawn("ssh %s -l %s %s" % (self.ssh_options, username, host))
        #self.child.logfile = sys.stderr

        index = 0
        pass_cnt = 0
        while 1:
            if (index > 2):
                raise LoginError("Login loop. Possibly bad prompt")

            result = self.child.expect(self.LOGIN_EXPECTS + [self.prompt], timeout=self.login_timeout)
    
            #Password needed
            if (result == self.MATCH_PASSWORD):
                if not (pass_cnt > 0):
                    self.child.sendline(password) 
                else: 
                    self.child.close()
                    raise AuthenticationError("Bad password")
                pass_cnt = pass_cnt + 1
    
            #Need to accept the client key
            elif (result == self.MATCH_ACCEPT_LICENSE):
                self.child.sendline("yes")

            #Returned permission denied
            elif (result == self.MATCH_PERM_DENIED):
                self.child.close()
                raise AuthenticationError("Permission denied")

            #Returned permission denied
            elif (result == self.MATCH_NO_ROUTE):
                self.child.close()
                raise NetworkError("No route to host")

            #Returned connection refused
            elif (result == self.MATCH_CONN_REFUSED):
                self.child.close()
                raise NetworkError("Connection refused")

            #Returned connection refused
            elif (result == self.MATCH_EOF):
                self.child.close()
                raise LoginError("Connection failed without an error")

            #Returned connection refused
            elif (result == self.MATCH_TIMEOUT):
                self.child.close()
                raise LoginError("Connection timed out looking for prompt/error")
  
            # Are we in?
            elif (result == self.MATCH_PROMPT):
                self.last_prompt = self.child.after
                break;

            index = index + 1

  
class TelnetTransport(BaseTransport):
    def __init__(self, login_timeout=DEFAULT_LOGIN_TIMEOUT, 
                       *args, **kwargs):
        super(TelnetTransport, self).__init__(*args, **kwargs)
        self.port = 23
        self.login_timeout=login_timeout

    LOGIN_EXPECTS = ["(?i)password",					# Password:    		     0
		             "(?i)username",					# Username 	 	         1 
                     "(?i)no route to host",            # No route to host       2
                     "(?i)connection refused",          # Connection refused     3
                     pexpect.EOF,						# Returned without error 4
                     pexpect.TIMEOUT,					# We timed out           5  
                    ]
    MATCH_PASSWORD       = 0
    MATCH_USERNAME       = 1
    MATCH_NO_ROUTE       = 2
    MATCH_CONN_REFUSED   = 3
    MATCH_EOF            = 4
    MATCH_TIMEOUT        = 5
    MATCH_PROMPT         = len(LOGIN_EXPECTS)

    def open(self, host, username="", password="", port=23):
        self.child = pexpect.spawn("telnet %s" % host)
        #self.child.logfile = sys.stderr
   
        index = 0
        pass_cnt = 0
        user_cnt = 0
        while 1:
            if (index > 2):
                raise LoginError("Login loop. Possibly bad prompt")

            result = self.child.expect(self.LOGIN_EXPECTS + [self.prompt], timeout=self.login_timeout)

            #Password needed
            if (result == self.MATCH_PASSWORD):
                if not (pass_cnt > 0):
                    self.child.sendline(password) 
                else: 
                    # Got password field twice
                    self.child.close()
                    raise AuthenticationError("Bad password")
                pass_cnt = pass_cnt + 1
      
            #Username
            elif (result == self.MATCH_USERNAME):
                if not (user_cnt > 0):
                    self.child.sendline(username)
                else:
                    # Got username twice
                    self.child.close()
                    raise AuthenticationErrorssc	("Failed to login:")
                user_cnt = user_cnt + 1

            #Returned permission denied
            elif (result == self.MATCH_NO_ROUTE):
                self.child.close()
                raise LoginError("No route to host")

            #Returned connection refused
            elif (result == self.MATCH_CONN_REFUSED):
                self.child.close()
                raise NetworkError("Connection refused")

            #Returned connection refused
            elif (result == self.MATCH_EOF):
                self.child.close()
                raise NetworkError("Connection failed without an error")

            #Returned connection refused
            elif (result == self.MATCH_TIMEOUT):
                self.child.close()
                raise LoginError("Connection timed out looking for prompt/error")

            # Are we in?
            elif (result == self.MATCH_PROMPT):
                self.last_prompt = self.child.after
                break;
